open Base
open OUnit2
open Helper

module Logging_graphics : Ax.React_loop.GRAPHICS_LAYER = struct
  let clear () = log "clear"
  let rect aabb color =
    logf "rect{%s;%s}"
      (Ax.Math.AABB.to_string aabb)
      (Ax.Color.to_string color)
end

let paths_to_lists =
  Sequence.map ~f:(fun (p,u) ->
      (Ax.Draw_tree.Path.to_list p, u))

;;
run_test_tt_main @@
test_list
  [
    ts (module Int)
      ["math1", 3 + 4, 7;
       "math2", 8 * 2, 16];
    ts (module IntOption)
      ["opt1", Some(3), Some(1 + 2);
       "opt2", None,    None];
    t_log "hello_world" ["hello";"world"]
      (fun () -> log "hello"; logf "w%sld" "or");
    tfails
      [ "fail_world", Failure("ello"),
        (fun () -> failwith "hello world")
      ; "math3", Division_by_zero,
        (fun () -> 3 / (2 - 4 / 2)) ];
    ts (module ListUnordered_(Int))
      ["uo1", [1;2;3], [3;2;1];
       "uo2", [1;3;4;3], [3;3;1;4];
       "uo3", [], []];

    "Draw_tree">:::begin
      let open Ax.Color in
      let open Ax.Math in
      let open Ax.Draw_tree in
      let t_render name (expected : Draw.t list) (tree : tree) =
        t name (module List_(Draw))
          expected
          (Sequence.to_list (render tree))
      in
      let rects = [ figure (Rect({w=30;h=40}, red))
                  ; figure (Rect({w=70;h=20}, blue))
                  ; figure (Rect({w=10;h=50}, black)) ] in

      [ t_render "H_rects"
          [ Rect({bx=0;by=5;bw=30;bh=40},red)
          ; Rect({bx=30;by=15;bw=70;bh=20},blue)
          ; Rect({bx=100;by=0;bw=10;bh=50},black) ]
          (arm `H rects)

      ; t_render "V_rects"
          [ Rect({bx=20;by=0;bw=30;bh=40},red)
          ; Rect({bx=0;by=40;bw=70;bh=20},blue)
          ; Rect({bx=30;by=60;bw=10;bh=50},black) ]
          (arm `V rects)

      ; t_render "H+V"
          [ Rect({bx=0;by=35;bw=20;bh=20},red)
          ; Rect({bx=25;by=0;bw=30;bh=40},blue)
          ; Rect({bx=20;by=40;bw=40;bh=50},green)
          ; Rect({bx=60;by=30;bw=10;bh=10},yellow)
          ; Rect({bx=60;by=40;bw=10;bh=20},cyan) ]
          (arm `H [ figure (Rect({w=20;h=20}, red))
                  ; arm `V [ figure (Rect({w=30;h=40}, blue))
                           ; figure (Rect({w=40;h=50}, green)) ]
                  ; arm `V [ figure (Rect({w=10;h=10}, yellow))
                           ; figure (Rect({w=10;h=20}, cyan)) ] ])
      ]
    end;

    "Draw_tree.capture_point">:begin
      let open Ax.Color in
      let open Ax.Math in
      let open Ax.Draw_tree in
      let id1, id2 = Ax.Event.Id.Private.(next_id ~diff:0, next_id ~diff:1) in
      let r1 = figure (Rect({w=10;h=10},red)) in
      let r2 = figure (Rect({w=10;h=10},red)) in
      let cap x y t = capture_point ~pos:{x;y} t in
      ts (module Option_(Ax.Event.Id))
        [ "empty",
          None,         cap 5 5 @@ arm `H []
        ; "didnt_hit",
          None,         cap 35 5 @@ arm `H [ r1 ; r2 ]
        ; "did_hit",
          Some(id1),    cap 5 5 @@ arm `H [ r1 ; r2 ]
        ; "did_hit'",
          Some(id2),    cap 15 5 @@ arm `H [ r1 ; r2 ]
        ]
    end;

    "Generic_tree.Zipper">:begin
      let open Ax.Private.Generic_tree in
      let module Z = Ax.Private.Generic_tree.Zipper in
      let module Tree = struct
        type nonrec t = (char, string) t
        include Make_show_eq(Char)(String)
      end in
      let lf_x = Leaf("xx") in let lf_y = Leaf("yy") in let lf_z = Leaf("zz") in
      let build mid = Arm('P', [ lf_x ; Arm('Q', [ lf_y ; lf_z ; mid ]) ; lf_z ; lf_y ]) in
      let zip = Z.(build lf_x |> of_tree |> down_exn 1 |> down_exn 2) in
      ts (module Tree)
        [ "get_focus",
          lf_x,
          Z.get zip
        ; "set_focus",
          build lf_y,
          Z.(zip |> set lf_y |> to_tree)
        ; "set_up",
          Arm('P', [lf_x;lf_y;lf_z;lf_y]),
          Z.(zip |> up_exn |> set lf_y |> to_tree)
        ; "side_zero",
          build lf_x,
          Z.(zip |> side_exn 0 |> to_tree)
        ; "side_positive",
          Arm('P', [ lf_x ; Arm('Q', [lf_y;lf_z;lf_z]) ; lf_z ; lf_y ]),
          Z.(of_tree (build lf_x) |> down_exn 1 |> down_exn 0 |> side_exn 2 |> set lf_z |> to_tree)
        ; "side_negative",
          Arm('P', [ lf_x ; Arm('Q', [lf_x;lf_z;lf_x]) ; lf_z ; lf_y ]),
          Z.(of_tree (build lf_x) |> down_exn 1 |> down_exn 2 |> side_exn ~-2 |> set lf_x |> to_tree)
        ]
    end;

    "Draw_tree.update">:::begin
      let open Ax.Color in
      let open Ax.Draw_tree in
      let open Ax.Draw_tree.Update in
      let f1 = Rect({w=1;h=1}, red) in
      let f2 = Rect({w=2;h=2}, blue) in
      let f3 = Rect({w=3;h=3}, white) in
      let f4 = Rect({w=4;h=4}, black) in
      let f5 = Rect({w=5;h=5}, red) in
      let r1 = figure f1 in let r2 = figure f2 in
      let r3 = figure f3 in let r4 = figure f4 in let r5 = figure f5 in
      let t_orig = arm `H [ r1 ; arm `V [ r2 ; r3 ; arm `H [r4;r5] ] ] in
      let t1     = arm `H [ r4 ; arm `V [ r2 ; r3 ; arm `H [r4;r5] ] ] in
      let t2     = arm `H [ r1 ; arm `V [ r2 ; r3 ; arm `H [r1;r5] ] ] in
      let t3     = arm `H [ r1 ; arm `V [ r2 ; r5 ; arm `H [r1;r5] ] ] in
      let t4     = arm `H [ r1 ; arm `V [ r2 ; r3 ; arm `H [r1;r2] ] ] in
      let apply_updates ?(targ=t_orig) us =
        apply_updates_exn
          (Sequence.map (Sequence.of_list us)
             ~f:(fun (l, fg) -> (Path.of_list l, fg)))
          targ
      in

      [ ts (module Ax.Draw_tree) @@
        List.map ~f:(fun (n,t',us) -> (n, t', apply_updates us))
          [ "fig_shallow",  t1, [ [0],     Set(f4)                    ]
          ; "fig_deep",     t2, [ [1;2;0], Set(f1)                    ]
          ; "fig_two",      t3, [ [1;1],   Set(f5) ; [1;2;0], Set(f1) ]
          ; "fig_two'",     t3, [ [1;2;0], Set(f1) ; [1;1],   Set(f5) ]
          ; "fig_adjacent", t4, [ [1;2;0], Set(f1) ; [1;2;1], Set(f2) ] ]
      ; tfail "fig_replace_arm"
          (Failure("cannot update non-figure"))
          (fun () -> apply_updates [ [1], Set(f1) ])

      ; t "mount_uninit" (module Ax.Draw_tree)
          t3
          (apply_updates ~targ:uninit
             [ [], Mount(t3) ])
      ; t "mount_uninit'" (module Ax.Draw_tree)
          (arm `H [ r1 ; r2 ])
          (apply_updates ~targ:(arm `H [ uninit ; r2 ])
             [ [0], Mount(r1)])
      ; t "mount_update_path" (module Ax.Draw_tree.Path)
          (Path.of_list [0;1])
          (let t = figure f1 in
           let path = Ax.Draw_tree.path t in
           let _ = apply_updates ~targ:(arm `H [ arm `V [ r1 ; uninit ]])
                     [ [0;1], Mount(t) ] in
           path)
      ; t "unmount" (module Ax.Draw_tree)
          (arm `H [ uninit ; r2 ; uninit ])
          (apply_updates ~targ:(arm `H [r1;r2;r3])
             [ [0], Unmount
             ; [2], Unmount ])

      ; tfail "mount_init"
          (Failure("cannot mount already-initialized"))
          (fun () -> apply_updates [ [], Mount(t1) ])
      ; tfail "mount_init'"
          (Failure("cannot mount already-initialized"))
          (fun () -> apply_updates [ [1;2;0], Mount(t1) ])
      ]
    end;

    "Source">:::begin
      let open Ax.Source in
      let src1 = create () in
      let src2 = create () in
      [ tfail "empty_state" (Failure("source unset"))
          (fun () -> State.(get_exn src1 empty))
      ; t "not_equal" (module Bool)
          false
          (equal src1 src2)
      ; t "set_get" (module Int)
          5
          State.(get_exn src1 (set src1 5 empty))
      ; t "set_set_get" (module Int)
          8
          State.(get_exn src1 (set src1 8 (set src1 5 empty)))
      ; t "set_other_get" (module Int)
          5
          State.(get_exn src1 (set src2 8 (set src1 5 empty)))
      ]
    end;

    "Component">:::begin
      let module T = Ax.Draw_tree in
      let module O = Ax.Component in
      let module B = Ax.Behavior in
      let module S = Ax.Source in
      let open Ax.Color in
      let open Ax.Math in
      let rect1 = T.Rect({w=1;h=1}, blue) in
      let rect2 = T.Rect({w=2;h=2}, blue) in
      let rect3 = T.Rect({w=3;h=3}, white) in
      let dyn s f = O.Dynamic(B.(map (of_source s) ~f)) in
      let dyn_rect b = O.(rect B.(b >>| fun t -> ({w=t;h=t}, blue))) in

      let rec_subsink_ids sink_id (st : O.state) =
        O.recursive_subsink_ids
          (snd (Map.find_exn st.sinks sink_id))
          st
      in

      [
        "mount">:::begin
          let st0 = O.empty_state (S.State.of_time 2) in
          let t name exp_tree com = t name (module T) exp_tree (snd (O.mount com st0)) in
          [ t "crect"
              T.(figure rect1)
              O.(crect {w=1;h=1} blue)
          ; t "time_rect"
              T.(figure rect2)
              (dyn_rect B.time)
          ; t "arm"
              T.(arm `H [ figure rect1 ; figure rect2 ])
              O.(arm `H [ crect {w=1;h=1} blue ; dyn_rect B.time ])
          ; t "fold"
              T.(figure rect2)
              O.(fold ~init:2 ~f:(fun _ x -> x)
                   dyn_rect) ]
        end;

        "updates">:::begin
          let other_src = S.create () in
          let irrelevant_src = S.create () in
          let (st, _tree) =
            O.(init (S.State.(of_time 2 |> set other_src 3))
                 (arm `H [ crect {w=1;h=1} blue
                         ; dyn_rect B.time
                         ; rect B.(of_source other_src >>|
                                   fun v -> ({w=v;h=v}, blue)) ]))
          in
          let t name exp act =
            t name (module ListUnordered_(Pair_(List_(Int))(T.Update)))
              exp (Sequence.to_list (paths_to_lists act))
          in
          let change src value =
            snd (O.source_changed src value st)
          in
          [ t "change_time"
              [ [1], Unmount
              ; [1], Mount(T.(figure (Rect({w=3;h=3}, blue)))) ]
              (change S.time 3)
          ; t "change_other"
              [ [2], Unmount
              ; [2], Mount(T.(figure (Rect({w=2;h=2}, blue)))) ]
              (change other_src 2)
          ; t "change_irr"
              []
              (change irrelevant_src 1) ]
        end;

        "state">:::begin
          let src1 = S.create () in
          let src2 = S.create () in
          let srcs0 = S.State.(empty |> set src1 1 |> set src2 2) in

          [ "static">:::
            (let st, _ = O.(init srcs0
                              (arm `V [ crect {w=1;h=1} blue
                                      ; crect {w=2;h=2} blue ])) in
             [ t "conns" (module Bool) true (Map.is_empty st.conns)
             ; t "sinks" (module Bool) true (Map.is_empty st.sinks) ])

          ; "parallel">:::
            (let sink_id1 = O.Sink_id.anticipate 0 in
             let sink_id2 = O.Sink_id.anticipate 1 in
             let st, _ = O.(init srcs0
                              (arm `V [ dyn_rect (B.of_source src1)
                                      ; dyn_rect (B.of_source src2) ])) in
             let st', _ = O.source_changed src1 99 st in
             [ t "conns" (module MultiMap_(S)(O.Sink_id))
                   (Map.of_alist_exn (module S)
                      [ src1, [ sink_id1 ];
                        src2, [ sink_id2 ] ])
                   st.conns
             ; t "sinks" (module Map_(O.Sink_id)(T.Path))
                 (Map.of_alist_exn (module O.Sink_id)
                    [ sink_id1, T.Path.of_list [0]
                    ; sink_id2, T.Path.of_list [1] ])
                 (Map.map ~f:fst st.sinks)
             ; t "sub_sinks" (module Map_(O.Sink_id)(ListUnordered_(O.Sink_id)))
                 (Map.of_alist_exn (module O.Sink_id)
                    [ sink_id1, []
                    ; sink_id2, [] ])
                 (Map.map ~f:(fun (_, s) -> s.sub_sinks) st.sinks)
             ; t "conns_after_update" (module MultiMap_(S)(O.Sink_id))
                 st.conns
                 st'.conns ])

          ; "nested">:::
            (let sink_id1 = O.Sink_id.anticipate 0 in
             let sink_id2 = O.Sink_id.anticipate 1 in
             let st, _ = O.(init srcs0
                              (dyn src1 (fun v1 ->
                                   if Int.(v1 = 1)
                                   then dyn_rect (B.of_source src2)
                                   else cfigure rect3))) in
             let st_1, _ = O.source_changed src1 99 st in
             let st_2, _ = O.source_changed src2 99 st in
             [ t "conns" (module MultiMap_(S)(O.Sink_id))
                   (Map.of_alist_exn (module S)
                      [ src1, [ sink_id1 ]
                      ; src2, [ sink_id2 ] ])
                   st.conns
             ; t "sinks" (module Map_(O.Sink_id)(T.Path))
                 (Map.of_alist_exn (module O.Sink_id)
                    [ sink_id1, T.Path.of_list []
                    ; sink_id2, T.Path.of_list [] ])
                 (Map.map ~f:fst st.sinks)
             ; t "sub_sinks" (module Map_(O.Sink_id)(ListUnordered_(O.Sink_id)))
                 (Map.of_alist_exn (module O.Sink_id)
                    [ sink_id1, [ sink_id2 ]
                    ; sink_id2, [] ])
                 (Map.map ~f:(fun (_, s) -> s.sub_sinks) st.sinks)
             ; t "rec_subsink_ids" (module ListUnordered_(O.Sink_id))
                 [ sink_id2 ]
                 (rec_subsink_ids sink_id1 st)
             ; t "conns_after_update/1" (module MultiMap_(S)(O.Sink_id))
                 (Map.of_alist_exn (module S)
                    [ src1, [ sink_id1 ];
                      src2, [ ] ])
                 st_1.conns
             ; t "rec_subsink_ids_after_update" (module ListUnordered_(O.Sink_id))
                 [ ]
                 (rec_subsink_ids sink_id1 st_1)
             ; t "conns_after_update/2" (module MultiMap_(S)(O.Sink_id))
                 st.conns
                 st_2.conns ])

          ; "very_nested">:::
            (let (sid1, sid2, sid3) =
               O.Sink_id.(anticipate 0, anticipate 1, anticipate 2)
             in
             let st, _ =
               O.(init srcs0 @@
                  dyn src1 @@ fun _ ->
                  dyn src2 @@ fun _ ->
                  dyn_rect (B.of_source src1))
             in
             [ t "conns" (module MultiMap_(S)(O.Sink_id))
                   (Map.of_alist_exn (module S)
                      [ src1, [ sid1 ; sid3 ]
                      ; src2, [ sid2 ] ])
                   st.conns
             ; t "sub_sinks" (module Map_(O.Sink_id)(ListUnordered_(O.Sink_id)))
                 (Map.of_alist_exn (module O.Sink_id)
                    [ sid1, [ sid2 ]
                    ; sid2, [ sid3 ]
                    ; sid3, [] ])
                 (Map.map ~f:(fun (_, s) -> s.sub_sinks) st.sinks)
             ; t "rec_subsink_ids" (module ListUnordered_(O.Sink_id))
                 [ sid2 ; sid3 ]
                 (rec_subsink_ids sid1 st) ])

          ; "fold">:::
            (let src = S.Private.next_id ~diff:0 in
             (* let sid = O.Sink_id.anticipate 0 in *)
             let st, _ =
               O.(init srcs0
                    (fold ~init:2 ~f:(fun _ -> (+) 1)
                       dyn_rect))
             in
             [ t "source" (module Int)
                 2
                 (S.State.get_exn src st.sources)
             ])
          ]
        end;
      ]
    end;

    (*
    "Component.mount">:::begin
      let module T = Ax.Draw_tree in
      let module C = Ax.Component in
      let module B = Ax.Behavior in
      let module S = Ax.Source in
      let module E = Ax.Event in
      let open Ax.Color in
      let open Ax.Math in
      let r1 = C.rect @@ B.const ({w=1;h=1}, red) in
      let r2 = C.rect @@ B.const ({w=2;h=2}, blue) in
      let r3 = C.rect @@ B.const ({w=3;h=3}, white) in
      let r_dyn b = C.rect @@ B.(b >>| fun t -> {w=t;h=t}, black) in
      let state = S.State.of_time 25 in

      let t_sinks name e a =
        t name (module ListUnordered_(Pair_(S)(T.Path)))
          e
          (List.bind (Map.to_alist a)
             ~f:(fun (src, sinks) ->
                 List.map sinks ~f:(fun (path, _) -> (src, path))))
      in

      [ t "static" (module Ax.Draw_tree)
          C.(let _, _, t =  mount state @@
               arm `H [ r1
                      ; arm `V [ r2
                               ; r3 ] ]
             in t)
          T.(arm `H [ figure (Rect({w=1;h=1}, red))
                    ; arm `V [ figure (Rect({w=2;h=2}, blue))
                             ; figure (Rect({w=3;h=3}, white)) ] ])

      ; "time">:::begin
          let _, callbacks, tree =
            C.(mount state @@
               arm `H
                 [ arm `V
                     [ r1
                     ; r_dyn B.time ]
                 ; r2
                 ])
          in
          [ t "init_tree" (module Ax.Draw_tree)
              T.(arm `H [ arm `V [ figure (Rect({w=1;h=1}, red))
                                 ; figure (Rect({w=25;h=25}, black)) ]
                        ; figure (Rect({w=2;h=2}, blue)) ])
              tree
          ; t_sinks "sinks"
              [ S.time, T.Path.of_list [0;1] ]
              callbacks.sinks
          ; t "handlers" (module ListUnordered_(E.Id))
              []
              (Map.keys callbacks.handlers) ]
        end

      ; "events">:::begin
          let src = S.Private.next_id ~diff:0 in
          let (eid1, eid2) = E.Id.Private.(next_id ~diff:1, next_id ~diff:2) in
          let init_state, callbacks, tree =
            C.(mount state @@
               arm `H
                 [ arm `V
                     [ r1
                     ; fold ~init:57 ~f:(fun n _ -> n + 1)
                         r_dyn ]
                 ; callback ~f:ignore r3
                 ])
          in
          [ t "init_tree" (module Ax.Draw_tree)
              T.(arm `H [ arm `V [ figure (Rect({w=1;h=1}, red))
                                 ; figure (Rect({w=57;h=57}, black)) ]
                        ; figure (Rect({w=3;h=3}, white)) ])
              tree
          ; t "init_state" (module Int)
              57 (S.State.get_exn src init_state)
          ; t_sinks "sinks"
              [ src, T.Path.of_list [0;1] ]
              callbacks.sinks
          ; t "handlers" (module ListUnordered_(E.Id))
              [eid1;eid2]
              (Map.keys callbacks.handlers)
          ]
        end;
      ]
    end;

    "React_loop">:::begin
      let open Ax.Color in
      let open Ax.Math in
      let open Ax.Component in
      let module B = Ax.Behavior in
      let current_cb = ref ignore in
      let current_time = ref 0 in
      let tick () = Int.incr current_time ; !current_cb Ax.React_loop.Tick in
      let click x y = !current_cb (Ax.React_loop.Click {x;y}) in
      let exit () = !current_cb Ax.React_loop.Exit in
      let module RL =
        Ax.React_loop.Make
          (struct
            type callback_id = unit
            let bind_callback f = current_cb := f
            let unregister () = current_cb := ignore ; logf "unreg"
            let poll_time_ms () = 1000 * !current_time
          end)
          (Logging_graphics)
      in
      [ t_log "no_callbacks"
          [ "clear" ; "rect{<0,0 4x4>;#0ff}" ; "unreg" ]
          (fun () -> RL.run (rect @@ B.const ({w=4;h=4}, cyan))
                   ; tick ()
                   ; tick ()
                   ; exit ())

      ; t_log "simple"
          [ "clear" ; "rect{<0,0 25x26>;#0ff}"
          ; "clear" ; "rect{<0,0 26x27>;#0ff}"
          ; "clear" ; "rect{<0,0 27x28>;#0ff}"
          ; "unreg" ]
          (fun () -> current_time := 25
                   ; RL.run (rect B.(time >>| fun t -> {w=t;h=t+1}, cyan))
                   ; tick ()
                   ; tick ()
                   ; exit ())

      ; t_log "click"
          [ "clear" ; "rect{<0,0 25x30>;#0ff}"
          ; "CLICK <5,5>" ; "CLICK <15,5>"
          ; "unreg" ]
          (fun () ->
             let f = function Ax.Event.Click(pos) ->
               logf "CLICK %s" (Pos.to_string pos)
             in
             RL.run (callback ~f (rect (B.const ({w=25;h=30}, cyan))))
           ; click 5 5 ; click 28 5 ; click 15 5 ; click 5 37
           ; exit ())

      ; t_log "click+state"
          [ "clear" ; "rect{<0,0 25x30>;#0ff}"
          ; "clear" ; "rect{<0,0 26x31>;#0ff}"
          ; "clear" ; "rect{<0,0 27x32>;#0ff}"
          ; "unreg" ]
          (fun () ->
             let inc n _ = n + 1 in
             let comp n = rect B.(n >>| fun n -> ({w=n;h=n+5}, cyan)) in
             RL.run (fold ~init:25 ~f:inc comp)
           ; click 5 5 ; click 28 5 ; click 15 5 ; click 5 37
           ; exit ())

      ]
    end;
*)

  ]
